///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/line_search.cpp
///
/// \brief   implementation file for the line search function
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<vector>
#include<numeric>
#include<cassert>
#include<algorithm>
#include<cmath>
#include<string>

#include<boost/format.hpp>
#include<boost/shared_ptr.hpp>
#include<boost/math/special_functions/fpclassify.hpp>

#include<formic/fqmc/line_search.h>
#include<formic/wfn/coeff_obj.h>
#include<formic/wfn/configuration.h>
#include<formic/fqmc/energy.h>
#include<formic/fqmc/sampler.h>
#include<formic/input/base.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
#include<formic/random/random.h>
#include<formic/output/output.h>
#include<formic/exception.h>
#include<formic/lapack/interface.h>
#include<formic/timing/timing.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  Perform a line search to minimize the energy along the provided update direction.
///
/// \param[in]      userinp        the user's input options
/// \param[in]      update         the update direction
/// \param[in]      max_sdev       maximum standard deviation allowed in the updates' energies
/// \param[in]      sampler        object that controls sampling
/// \param[in,out]  vf             a shared pointer to the value function
///                                On exit, the best update has been added to the value function
/// \param[in,out]  step_arg       on entry, the initial step size to use.
///                                on exit, the step size that was used.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::fqmc::line_search(const formic::InputBase & userinp,
                                                 const S * const update,
                                                 const double max_sdev,
                                                 boost::shared_ptr<formic::fqmc::Sampler<S> > & sampler,
                                                 const boost::shared_ptr<formic::CoeffObj<S> > & vf,
                                                 double & step_arg) {

  // start a timer
  formic::start_timer("line search");

  // print a greeting
  if (formic::mpi::rank() == 0)
    formic::of << boost::format("Starting a line search.") << std::endl << std::endl;

  // get the initial energy and it's standard deviation
  double sdev = 0.0;
  double serr = 0.0;
  const double init_energy = formic::fqmc::energy(userinp, sampler, vf, sdev, serr, false);
  if (formic::mpi::rank() == 0)
    formic::of << boost::format("step size = %16.8e         energy = %23.15f         sdev = %.2e         %s")
                  % 0.0
                  % init_energy
                  % sdev
                  % "initial energy"
               << std::endl;

  // initialize line search variables
  std::vector<double> step_sizes(1, 0.0);
  std::vector<double> step_energies(1, init_energy);
  double best_step_size = 0.0;
  double last_step_size = 0.0;
  double best_step_energy = init_energy;
  int best_step = 0;
  int step_i = 0;
  sdev = 0.0;
  serr = 0.0;
  int n_refine = 0;

  // save the current variables
  std::vector<S> saved_vars(vf->n_der(), formic::zero(S()));
  vf->to_vector(&saved_vars.at(0));

  // initialize length increase parameter for longer steps
  double longer_step_multiplier = 1.3;

  for (int snum = 0; snum < userinp.get<int>("line_search_steps"); snum++) {

    std::string message;

    double new_step_size = -1.0;

    // choose what to do next
//    if ( best_step != 0 && sdev > max_sdev && last_step_size < *step_sizes.rbegin() ) {
//      if (formic::mpi::rank() == 0)
//        formic::of << boost::format("stopping due to unexpectedly large standard deviation") << std::endl;
//      break;
//    }
    if ( snum == 0 ) {
      new_step_size = 1.0 * step_arg;
      message = ( boost::format("initial step") ).str();
    } else if ( sdev > max_sdev ) {
      if (best_step == 0) {
        new_step_size = step_sizes.at(1) * 0.5;
      } else {
        new_step_size = step_sizes.at(step_i) / longer_step_multiplier;
        longer_step_multiplier = 1.0 + 0.5 * ( longer_step_multiplier - 1.0 );
        new_step_size *= longer_step_multiplier;
      }
      message = ( boost::format("shorter step (large sdev)") ).str();
      step_sizes.erase(step_sizes.begin()+step_i);
      step_energies.erase(step_energies.begin()+step_i);
      if ( best_step == step_i )
        throw formic::Exception("unexpected deletion of best step due to large standard deviation in formic::fqmc::line_search");
      else if ( best_step > step_i )
        best_step--;
    //} else if ( snum < 1 ) {
    //  //new_step_size = step_arg;
    //  new_step_size = 0.5 * step_arg;
    //  message = ( boost::format("step of 0.5 * step_arg") ).str();
    //} else if ( snum < 3 && best_step == 2 ) {
    //  //new_step_size = step_arg;
    //  new_step_size = 2.0 * step_arg;
    //  message = ( boost::format("step of 2.0 * step_arg") ).str();
    //} else if ( step_sizes.size() < 5 ) {
    //  new_step_size = std::max(step_sizes.at(1), best_step_size) * ( best_step_energy == init_energy ? 0.5 : 1.3 );
    //  //new_step_size = step_sizes.at(step_i) * ( best_step_energy == init_energy ? 0.5 : 1.3 );
    //  message = ( boost::format("trying 4th step...") ).str();
    } else if ( best_step == 0 ) {
      new_step_size = step_sizes.at(1) * 0.5;
      message = ( boost::format("shorter step") ).str();
    } else if ( best_step == step_sizes.size()-1 ) {
      new_step_size = best_step_size * longer_step_multiplier;
      message = ( boost::format("longer step") ).str();
    } else {
      n_refine++;
      new_step_size = formic::porabola_min_max(&step_sizes.at(best_step-1), &step_energies.at(best_step-1));
      // make sure the new step size is well behaved
      if ( new_step_size < 0.0 || boost::math::isnan(new_step_size) ) {
        break; // stop if we have found a negative number or a nan
      } else if ( boost::math::isinf(new_step_size) || new_step_size > 1.4 * (*step_sizes.rbegin()) ) {
        new_step_size = 1.4 * (*step_sizes.rbegin());
      } else if ( new_step_size < 0.6 * step_sizes.at(1) ) {
        new_step_size = 0.6 * step_sizes.at(1);
      }
      message = ( boost::format("refining minimum") ).str();
    }
    formic::mpi::bcast(new_step_size);

    // prepare the next step
    for (step_i = 0; step_i < step_sizes.size(); step_i++) {
      if ( new_step_size < step_sizes.at(step_i) )
        break;
    }
    step_sizes.insert(step_sizes.begin()+step_i, new_step_size);
    step_energies.insert(step_energies.begin()+step_i, 1.0e100);
    if ( step_i <= best_step )
      best_step++;

    // compute new energy and standard deviation using correlated sampling
    vf->update_variables(update,  formic::unity(S()) * step_sizes.at(step_i));
    step_energies.at(step_i) = formic::fqmc::energy(userinp, sampler, vf, sdev, serr, false);
    vf->from_vector(&saved_vars.at(0));
    if (formic::mpi::rank() == 0)
      formic::of << boost::format("step size = %16.8e         energy = %23.15f         stat error = %.2e         %s")
                    % step_sizes.at(step_i)
                    % step_energies.at(step_i)
                    % serr
                    % message
                 << std::endl;

    // if this was the best step size so far, save it
    if ( sdev < max_sdev && step_energies.at(step_i) < best_step_energy ) {
      best_step = step_i;
      best_step_size = step_sizes.at(step_i);
      best_step_energy = step_energies.at(step_i);
    }

    // remember what the last step size was
    last_step_size = new_step_size;

    // if we have done three minima refinements, stop searching
    if (n_refine > 2)
      break;

    // if the step sizes have become too close together, stop searching
    const double step_thresh = 1.0e-6;
    if ( best_step < step_sizes.size()-1 ) {
      if ( std::abs( ( step_sizes.at(best_step) - step_sizes.at(best_step+1) ) / step_sizes.at(best_step+1) ) < step_thresh )
        break;
    }
    if ( best_step > 0 ) {
      if ( std::abs( ( step_sizes.at(best_step) - step_sizes.at(best_step-1) ) / step_sizes.at(best_step) ) < step_thresh )
        break;
    }

  }

  // if a zero step was best, do not apply any update
  if ( best_step == 0 ) {

    if (formic::mpi::rank() == 0)
      formic::of << std::endl << "The line search failed to find a step that lowers the energy.  No update was applied." << std::endl << std::endl;

  // otherwise, apply the update corresponding to the best step size
  } else {

    vf->update_variables(update,  formic::unity(S()) * best_step_size);

    if (formic::mpi::rank() == 0)
      formic::of << std::endl
                 << boost::format("The line search's best step size was %16.8e with an energy of %23.15f") % best_step_size % best_step_energy
                 << std::endl
                 << std::endl
                 << boost::format("The corresponding update has been applied.")
                 << std::endl
                 << std::endl;

  }

  // stop the timer
  formic::stop_timer("line search");

  // return the step size used
  step_arg = best_step_size;

}

template void formic::fqmc::line_search(const formic::InputBase & userinp,
                                        const double * const update,
                                        const double max_sdev,
                                        boost::shared_ptr<formic::fqmc::Sampler<double> > & sampler,
                                        const boost::shared_ptr<formic::CoeffObj<double> > & vf,
                                        double & step_arg);
template void formic::fqmc::line_search(const formic::InputBase & userinp,
                                        const std::complex<double> * const update,
                                        const double max_sdev,
                                        boost::shared_ptr<formic::fqmc::Sampler<std::complex<double> > > & sampler,
                                        const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > & vf,
                                        double & step_arg);
